# database/postgres_manager.py
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import SQLAlchemyError
from contextlib import contextmanager
import pandas as pd
from typing import Iterator, Optional, Any
import os
import sys

# 添加项目根目录到Python路径
current_file = os.path.abspath(__file__)
stock_data_system_dir = os.path.dirname(os.path.dirname(current_file))
stock_traing_root = os.path.dirname(stock_data_system_dir)

if stock_traing_root not in sys.path:
    sys.path.insert(0, stock_traing_root)

from stock_data_system.config.database import DatabaseConfig
from stock_data_system.database.models import Base

class PostgresManager:
    """PostgreSQL数据库管理类"""

    def __init__(self, config: DatabaseConfig):
        self.config = config
        self.engine = self._create_engine()
        self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)

    def _create_engine(self):
        """创建数据库引擎"""
        connection_string = (
            f"postgresql://{self.config.username}:{self.config.password}"
            f"@{self.config.host}:{self.config.port}/{self.config.database}"
        )

        return create_engine(
            connection_string,
            pool_size=self.config.pool_size,
            max_overflow=self.config.max_overflow,
            echo=self.config.echo_sql
        )

    def create_tables(self):
        """创建所有表"""
        try:
            Base.metadata.create_all(bind=self.engine)
            print("数据库表创建成功")
        except SQLAlchemyError as e:
            print(f"创建表失败: {e}")
            raise

    def check_connection(self) -> bool:
        """检查数据库连接"""
        try:
            with self.engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            print("数据库连接成功")
            return True
        except SQLAlchemyError as e:
            print(f"数据库连接失败: {e}")
            return False

    @contextmanager
    def get_session(self) -> Iterator[Any]:
        """获取数据库会话"""
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def execute_query(self, query: str, params: dict = None) -> pd.DataFrame:
        """执行查询并返回DataFrame"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text(query), params or {})
                df = pd.DataFrame(result.fetchall(), columns=result.keys())
                return df
        except SQLAlchemyError as e:
            print(f"查询执行失败: {e}")
            return pd.DataFrame()

    def bulk_insert_dataframe(self, df: pd.DataFrame, table_name: str, if_exists: str = 'append'):
        """批量插入DataFrame数据"""
        if df.empty:
            print("DataFrame为空，跳过插入")
            return

        try:
            df.to_sql(
                table_name,
                self.engine,
                if_exists=if_exists,
                index=False,
                method='multi',
                chunksize=1000
            )
            print(f"成功插入 {len(df)} 条数据到 {table_name}")
        except Exception as e:
            print(f"批量插入失败: {e}")

    def get_last_trade_date(self, symbol: str = None) -> Optional[str]:
        """获取最后交易日期"""
        try:
            if symbol:
                query = "SELECT MAX(trade_date) FROM daily_quotes WHERE symbol = :symbol"
                params = {'symbol': symbol}
            else:
                query = "SELECT MAX(trade_date) FROM daily_quotes"
                params = {}

            result = self.execute_query(query, params)
            if not result.empty and result.iloc[0, 0] is not None:
                return result.iloc[0, 0].strftime('%Y-%m-%d')
            return None
        except Exception as e:
            print(f"获取最后交易日期失败: {e}")
            return None

    def execute_raw_sql(self, sql: str, params: dict = None):
        """执行原始SQL语句"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text(sql), params or {})
                conn.commit()
                return result
        except SQLAlchemyError as e:
            print(f"SQL执行失败: {e}")